import torch
import glob
import os
from torchvision import transforms
import cv2
from PIL import Image
import pandas as pd
import numpy as np
from imgaug import augmenters as iaa
from utils import get_label_info, one_hot_it, RandomCrop, reverse_one_hot, one_hot_it_v11, one_hot_it_v11_dice
import random

def augmentation():
    # augment images with spatial transformation: Flip, Affine, Rotation, etc...
    # see https://github.com/aleju/imgaug for more details
    pass
def augmentation_pixel():
    # augment images with pixel intensity transformation: GaussianBlur, Multiply, etc...
    pass
import glob
class CamVid(torch.utils.data.Dataset):
    def __init__(self, image_path, label_path, csv_path, scale, loss='dice', mode='train'):
        super().__init__()
        self.mode = mode
        print(mode,label_path,end="")
        self.image_list=[]
        asd=glob.glob(label_path) #[:2000]
        if mode=='train':
            for i in range(len(asd)):
                if i%10!=2:
                    self.image_list.append(asd[i])
        else:
            for i in range(len(asd)):
                if i%10==2:
                    self.image_list.append(asd[i])
        a=cv2.imread(self.image_list[0])
        print(a.shape)
        print(f"imglen:{len(self.image_list)}")
        self.label_list = [v[:-3]+"png"  for v in self.image_list]
        self.fliplr = iaa.Fliplr(0.5)
        self.label_info = get_label_info(csv_path)
        # resize
        # self.resize_label = transforms.Resize(scale, Image.NEAREST)
        # self.resize_img = transforms.Resize(scale, Image.BILINEAR)
        # normalization
        self.to_tensor = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
            ])
        # self.crop = transforms.RandomCrop(scale, pad_if_needed=True)
        self.image_size = scale
        self.scale = [0.5, 1, 1.25, 1.5, 1.75, 2]
        self.loss = loss
        self.index=0
    def __getitem__(self, index):
        # load image and crop
        seed = random.random()
        img = Image.open(self.image_list[index]).resize((512,512)).convert("RGB")

        # if self.index==0:
        #     img.save("img.jpg")
        # random crop image
        # =====================================
        # w,h = img.size
        # th, tw = self.scale
        # i = random.randint(0, h - th)
        # j = random.randint(0, w - tw)
        # img = F.crop(img, i, j, th, tw)
        # =====================================

        scale = random.choice(self.scale)
        scale = (int(self.image_size[0] * scale), int(self.image_size[1] * scale))

        # randomly resize image and random crop
        # =====================================
        if self.mode == 'train':
            img = transforms.Resize(scale, Image.BILINEAR)(img)
            img = RandomCrop(self.image_size, seed, pad_if_needed=True)(img)
        # =====================================

        img = np.array(img)
        # load label
        label =Image.open(self.label_list[index]).resize((512,512)) # cv2.imread(self.label_list[index],-1)
        print(label)
        print(np.max(label))
        # x = np.dstack([x, x, x])
        # rgb=cv2.merge([x,x,x])

        # cv2.imwrite("asd.png",rgb)
        # label = Image.open("asd.png")
        # print(label)
        # label=Image.fromarray(rgb, 'RGB').resize((512,512))
        # label = Image.open(self.label_list[index])
        # if self.index==0:
        #     label.save("label.jpg")

        # crop the corresponding label
        # =====================================
        # label = F.crop(label, i, j, th, tw)
        # =====================================

        # randomly resize label and random crop
        # =====================================
        if self.mode == 'train':
            label = transforms.Resize(scale, Image.NEAREST)(label)
            label = RandomCrop(self.image_size, seed, pad_if_needed=True)(label)
        # =====================================

        label = np.array(label)

        # if self.mode!='train':
        #     # print(label.shape)
        #     label=cv2.resize(label,(512,512),cv2.INTER_NEAREST)
        # augment image and label
        if self.mode == 'train':
            seq_det = self.fliplr.to_deterministic()
            img = seq_det.augment_image(img)
            label = seq_det.augment_image(label)


        # image -> [C, H, W]
        img = Image.fromarray(img)
        img = self.to_tensor(img).float()
        if self.loss == 'dice':
            # label -> [num_classes, H, W]
            label = one_hot_it_v11_dice(label, self.label_info).astype(np.uint8)

            label = np.transpose(label, [2, 0, 1]).astype(np.float32)
            # print(label.shape)
            # label = label.astype(np.float32)
            label = torch.from_numpy(label)

            return img, label

        elif self.loss == 'crossentropy':
            label = one_hot_it_v11(label, self.label_info).astype(np.uint8)
            # label = label.astype(np.float32)
            label = torch.from_numpy(label).long()
        if self.index<self.image_list:
            self.index+=1
        else:
            self.index=0
            return img, label
    # def debug(i):
        # img = Image.open(self.image_list[index]).resize((960,640))
        # img.save("img.jpg")
        # x = cv2.imread(self.label_list[index],-1)*200
        # rgb=cv2.merge([x,x,x])
        # label=Image.fromarray(rgb, 'RGB').resize((960,640))
        # label.save("label.jpg")
    def __len__(self):
        return len(self.image_list)


if __name__ == '__main__':
    # data = CamVid('/path/to/CamVid/train', '/path/to/CamVid/train_labels', '/path/to/CamVid/class_dict.csv', (640, 640))
    data = CamVid(['./dataset/CamVid/train', './dataset/CamVid/val'],
                  ['./dataset/CamVid/train_labels', './dataset/CamVid/val_labels'], './dataset/CamVid/class_dict.csv',
                  (720, 960), loss='crossentropy', mode='val')
    from model.build_BiSeNet import BiSeNet
    from utils import reverse_one_hot, get_label_info, colour_code_segmentation, compute_global_accuracy

    label_info = get_label_info('/data/sqy/CamVid/class_dict.csv')
    for i, (img, label) in enumerate(data):
        print(label.size())
        print(torch.max(label))

